tensor([[0.5800, 0.5800, 0.5800, 0.5800, 0.5800],
        [0.4200, 0.4200, 0.4200, 0.4200, 0.4200]], grad_fn=<RoundBackward1>)